/*
  Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

// $Id$

#include<iostream>
#include "CCfits/CCfits"
#include "counter.h"
#include "polychromatic_grid.h"
#include "optical.h"
#include "scatterer.h"
#include "constants.h"
#include "dust_model.h"
#include "emission.h"
#include "emission-fits.h"
#include "full_sed_emergence.h"
#include "full_sed_grid.h"
#include "emergence-fits.h"
#include "mcrx-units.h"
#include "grid_factory.h"
#include "grid-fits.h"
#include "shoot.h"
#include "fits-utilities.h"
#include "boost/lexical_cast.hpp"
#include "boost/shared_ptr.hpp"
#include "density_generator.h"
#include "preferences.h"
#include "model.h"
#include "blackbody_grain.h"
#include "molecule_grain_model.h"
#include "p04_grain_model.h"
#include "grain_model_hack.h"
#include "xfer.h"
#include "imops.h"
#include "ir_grid.h"
#include "equilibrium.h"
#include "boost/lambda/lambda.hpp"
#include "boost/program_options.hpp"
#include "scatterer.h"
#include "atmosphere.h"

using namespace blitz;
using namespace mcrx;
using namespace CCfits;
using namespace std; 
using namespace boost::lambda;
using boost::lexical_cast;
namespace po = boost::program_options;

typedef local_random T_rng_policy;
typedef scatterer<polychromatic_scatterer_policy, T_rng_policy> T_scatterer;
typedef T_scatterer::T_lambda T_lambda;
typedef T_scatterer::T_biaser T_biaser;
typedef dust_model<T_scatterer, cumulative_sampling, T_rng_policy> T_dust_model;
typedef full_sed_emergence T_emergence;

typedef full_sed_grid<adaptive_grid> T_grid;
typedef ir_grid<adaptive_grid, cumulative_sampling, local_random> T_ir_grid;
typedef emission<polychromatic_policy, T_rng_policy> T_emission;


class greenhouse_factory: public grid_factory<T_grid::T_data> {
public:
  const T_float h2o_rh_;
  const T_float rho_gr_;
  int maxlevel_, n_lambda_;
  atmosphere a_;

  greenhouse_factory (T_float wf, T_float rg, int ml, int nl) :
    h2o_rh_(wf), rho_gr_(rg), maxlevel_(ml), n_lambda_(nl) {};
  T_densities calc_rho(const vec3d& p) {
    // 2 species: blackbody ground and H2O atmo
    T_densities rho(2);
    //T_densities rho(1);
    // check if we are underground
    if(p[2]<0) {
      rho(0)=rho_gr_;
      rho(1)=0;
      return rho;
    }
    // get density from model atmo
    rho(0)=0;
    // approximate saturation pressure 
    const T_float t=a_.temperature(p[2]);
    const T_float P=a_.pressure(p[2]);
    const T_float psat = (t>0) ?
      610.78 *exp( (t-273.15) / ( t-273.15+ 238.3 ) *17.2694 ) :
      exp( -6140.4 / ( t ) + 28.916 );
    // partial pressure of water
    T_float pw = psat*h2o_rh_;
    // check that partial pressure isn't higher than the total pressure
    if(pw>P) pw=P;
    rho(1)=a_.density(p[2])*pw/P;
    //rho(1)=0;
    return rho;
  };

  virtual bool refine_cell_p (const T_cell& c, int level) {
    if(level>maxlevel_) return false;

    // if we contain the ground level, refine to max
    const T_float zmin=c.getmin()[2];
    const T_float zmax=c.getmax()[2];
    if (zmin*zmax<0) return true;
    if ( (zmax>-0.1) && (zmin<-0.1) ) return true;
    return false;
  };
  virtual bool unify_cell_p (T_grid::const_local_iterator begin,
			     T_grid::const_local_iterator end) {
    return false;
  };

  virtual T_data get_data (const T_cell& c) {
    return T_data ( calc_rho(c.getcenter()), vec3d(0,0,0), 
		    T_lambda (n_lambda_) );
  }
    
  virtual int n_threads () {return 1;};
  virtual int work_chunk_levels () {return 0;};
  virtual int estimated_levels_needed () {return 0;};
};


void greenhouse_test (po::variables_map opt)
{
  const int n_threads = opt["n_threads"].as<int>();
  const long n_rays =opt["n_rays"].as<long>();
  const long n_rays_ir =opt["n_rays_ir"].as<long>();
  const int n_rays_int = opt["n_rays_int"].as<long>();
  const vec3i grid_n (5,5,3010);
  const T_float grid_extent_xy=100e3;
  const T_float grid_min_z=-10;
  const T_float grid_max_z=30e3;

  const int n_forced=opt["n_forced"].as<int>();
  const int n_ir_forced=opt["n_ir_forced"].as<int>();
  const int maxlevel=opt["maxlevel"].as<int>();
  const T_float I_rr=opt["i_rr"].as<T_float>();
  const T_float sundist = 100*grid_max_z;
  const T_float cameradist = 2*grid_max_z;
  const string output_file(opt["output-file"].as<string>());
  const int seed=opt["seed"].as<int>();

  const T_float rho_gr = opt["rho_gr"].as<T_float>();
  const T_float h2o_frac = opt["h2o_rh"].as<T_float>();
  
  Preferences prefs;
  prefs.setValue("n_threads", n_threads);
  prefs.setValue("use_grain_temp_lookup", false);
  prefs.setValue("grain_data_directory", string("/home/patrik/dust_data/crosssections"));
  prefs.setValue("n_threads", n_threads);
  prefs.setValue("use_cuda", false);

  mcrx::seed(seed);
  T_unit_map units;
  units ["length"] = "m";
  units ["mass"] = "kg";
  units ["luminosity"] = "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  // create wavelengths
  array_1 lambda(logspace(1.2e-7,1e-3,50));

  vector<T_float> lambdas(lambda.begin(),lambda.end());
  const int lambda550  = lower_bound(lambdas.begin(), lambdas.end(),
				     0.551e-6) - lambdas.begin()-1;
  const int vis_ref_lambda = lower_bound(lambdas.begin(), lambdas.end(),
				     opt["vis_ref"].as<T_float>()) - lambdas.begin()-1;;
  const int ir_ref_lambda = lower_bound(lambdas.begin(), lambdas.end(),
					opt["ir_ref"].as<T_float>()) - lambdas.begin()-1;;
  
  // load grain model
  T_dust_model::T_scatterer_vector sv;
  sv.push_back(boost::shared_ptr<T_scatterer> 
	       (new blackbody_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(1e-2, units)));
  sv.push_back(boost::shared_ptr<T_scatterer> 
	       //(new p04_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>("/home/patrik/dust_data/crosssections", prefs, units)));
  //(new blackbody_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(1e-2, units)));
	       (new molecule_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>("water.fits", prefs, units)));
  T_dust_model model (sv);
  model.set_wavelength(lambda); // this also resamples the grain_model objects

  // build grid
  cout << "Building grid" << endl;
  greenhouse_factory factory (h2o_frac, rho_gr, maxlevel, 
			      opt["do_ir"].as<bool>() ? lambda.size() : 0);

  // first create the adaptive grid
  boost::shared_ptr<adaptive_grid<T_grid::T_data> > gg
    (new adaptive_grid<T_grid::T_data>  
     (vec3d(-grid_extent_xy, -grid_extent_xy, grid_min_z),
      vec3d(grid_extent_xy, grid_extent_xy, grid_max_z),
      grid_n, factory));

  // and then the full_sed_grid
  T_grid gr(gg, units);

  cout << "\nSetting up emission" << endl;

  // calculate luminosity necessary to get normal insolation
  const T_float r=sqrt(2.0)*grid_extent_xy;
  const T_float area=constants::pi*r*r;
  const T_float L_em=1000*area;
  // size of solar disk should be 0.5deg
  const T_float r_sun=sundist*0.5/180*constants::pi;
  const T_float div=atan((r-r_sun)/sundist);

  blackbody bb(5800,L_em);
  auto_ptr<T_emission> em;
  em.reset(new divergent_floodlight_emission<polychromatic_policy, T_rng_policy> (vec3d (.01, .01, sundist), vec3d(0,0,-1),r_sun,div,bb.emission(lambda)));

  cout << "setting up emergence" << endl;
  std::vector<std::pair<T_float, T_float> > cam_pos;
  cam_pos.push_back(std:: make_pair (12.5*constants::pi/180, 0.1) );
  cam_pos.push_back(std:: make_pair (42.5*constants::pi/180, 0.1) );
  cam_pos.push_back(std:: make_pair (77.5*constants::pi/180, 0.1) );
  auto_ptr<T_emergence> e(new T_emergence(cam_pos, cameradist,
					  r, 300));

  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, T_grid, T_rng_policy> x (gr, *em, *e, model, rng, I_rr, 0,n_forced,10.,false, opt["do_ir"].as<bool>(),0);

  // calculate optical depths
  {
    T_densities n(x.integrate_ray(vec3d(0,0,0), vec3d(0,0,1)));
    T_lambda tau(model.total_abs_length(n(tensor::j)));
    cout << tau << endl;
    const int l  = lower_bound(lambdas.begin(), lambdas.end(),
			       10e-6) - lambdas.begin()-1;
    cout << "Column density upward: " << n << endl;
    cout << "Optical depth of atmosphere at 10um: " << tau(l) << endl;
    n=x.integrate_ray(vec3d(0,0,0), vec3d(0,0,-1));
    tau=model.total_abs_length(n(tensor::j));
    cout << "Column density downward: " << n << endl;
    cout << "Optical depth of ground at 10um: " << tau(l) << endl;

  }

  cout << "shooting" << endl;
  long n_rays_actual = 0;
  dummy_terminator t;
  T_biaser b(vis_ref_lambda);
  shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays, 
  	 n_rays_actual, n_threads, true, n_threads, 0, "");
  T_float normalization = 1./n_rays_actual;
  
  {
    FITS output("!"+output_file, Write);
    e->write_images(output,normalization,units, "-STAR", false, false);

    // calculate integrated SEDs
    Table* iq_hdu=output.addTable("INTEGRATED_QUANTITIES",0);
    iq_hdu->addColumn(Tdouble, "lambda", 1,
		      units.get("L_lambda"));
    write(iq_hdu->column("lambda"),lambda,1);
    for (int i=0; i<cam_pos.size(); ++i) {
      integrate_image(output,
		      "CAMERA"+lexical_cast<string>(i)+"-STAR",
		      4*constants::pi*cameradist*cameradist,
		      "L_lambda"+lexical_cast<string>(i),
		      lambda,
		      "",
		      units,
		      dummy_terminator());
    }
  }

  cout << "Dust_model absorbed stellar luminosity: " 
       << sum(model.absorbed()*delta_quantity(lambda))*normalization << endl;

  if(!opt["do_ir"].as<bool>())
    return;

  // now we start the interesting stuff...
  // make dust mass vector
  array_2 dust_mass(gr.n_cells(), model.n_scatterers());
  array_2 star_intensity(gr.intensities().shape());

  int cc=0;
  const T_float area_factor = gr.area_factor();
  for(T_grid::const_iterator c=gr.begin(); c!=gr.end(); ++c, ++cc) {
    dust_mass(cc, Range::all()) = c->data()->densities()*c->volume();
    star_intensity(cc, Range::all()) =
      (gr.intensities()(cc,Range::all())*normalization/
       (4*constants::pi*c->volume()*area_factor));
  }

  // first create the adaptive grid
  boost::shared_ptr<adaptive_grid<T_ir_grid::T_data> > irgg
    (new adaptive_grid<T_ir_grid::T_data>  
     (vec3d(-grid_extent_xy, -grid_extent_xy, grid_min_z),
      vec3d(grid_extent_xy, grid_extent_xy, grid_max_z),
      grid_n, gr.get_structure()));

  // and then the full grid
  T_ir_grid irg(irgg, dust_mass, lambda, units);

  array_2 dust_intensity(gr.n_cells(), lambda.size());
  dust_intensity=0;

  T_biaser b2(ir_ref_lambda);
  scatter_shooter<T_biaser> s(b2);
 
  // Determine dust equilibrium
  if (opt["do_equilibrium"].as<bool>())
    determine_dust_equilibrium_intensities
      (model,
       gr,
       irg,
       lambda,
       star_intensity,
       dust_intensity,
       states,
       n_rays_int,
       n_threads,
       true,
       n_threads,
       opt["ir_tol"].as<T_float>(),
       I_rr,
       n_ir_forced,
       10,
       s,
       t,
       true);
 
  // now generate output
  std::vector<grain_model<polychromatic_scatterer_policy, 
    mcrx_rng_policy>*> grain_models;
  for(int i=0; i< model.n_scatterers(); ++i) {
    grain_models.push_back
      (dynamic_cast<grain_model<polychromatic_scatterer_policy, 
       mcrx_rng_policy>*>(&model.get_scatterer(i)));
  }

  e.reset(new T_emergence(cam_pos, cameradist,
			  r, 300));
  irg.calculate_SED(star_intensity + dust_intensity, grain_models,
		    t, n_threads);
  gr.reset();
  xfer<T_dust_model, T_grid, T_rng_policy> x2 (gr, irg, *e, model, rng, 1e-2, 0,0,10.,false,false);
  
  cout << "Shooting dust emission" << endl;

  n_rays_actual = 0;
  shoot (x2, s, t, states, n_rays_ir,
  	 n_rays_actual, n_threads, true, n_threads, 0, "");

  normalization = irg.total_emission_weight()/n_rays_actual;
  FITS output(output_file, Write);
  e->write_images(output,normalization,units, "-IR", false, false);

  // calculate integrated SEDs
  for (int i=0; i<cam_pos.size(); ++i) {
    integrate_image(output,
		    "CAMERA"+lexical_cast<string>(i)+"-IR",
		    4*constants::pi*cameradist*cameradist,
		    "L_lambda_ir"+lexical_cast<string>(i),
		    lambda,
		    "",
		    units,
		    dummy_terminator());
  }
  
  // get temperatures
  {
    const T_grid::T_cell_tracker c = gr.locate(vec3d(0,0,-.1));
    const int i = gr.get_cell_number(c);
    const T_float temp = dynamic_cast<blackbody_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>*>(sv[0].get())->calculate_temp(star_intensity(i,Range::all())+dust_intensity(i,Range::all()));
    cout << "Ground temperature is " << temp << endl;
  }
  /*
  {
    ofstream out("cell_seds.txt");
    for(T_ir_grid::iterator c=irg.begin(); c!=irg.end(); ++c)
      out << sqrt(dot(c->getcenter(),c->getcenter())) << '\n';
    out << "\n\n";
    for(T_ir_grid::iterator c=irg.begin(); c!=irg.end(); ++c) {
      for(int l=0; l<lambda.size(); ++l)
	out << lambda(l) << '\t' << c->data()->get_emission()(l) << '\n';
      out << "\n\n";
    }
  }
  */
  {
    ofstream out("gridstruct.txt");
    for(T_ir_grid::iterator c=irg.begin(); c!=irg.end(); ++c)
      out << c->getcenter()[0] << '\t' << c->getcenter()[1] << '\t'
	  << c->getcenter()[2] << '\n';
  }
}


int main (int argc, char** argv)
{
  // Declare the supported options.
  po::options_description desc("IRDC model. Allowed options");
  desc.add_options()
    ("output-file", po::value<string>(), "output file name")
    ("n_threads", po::value<int>()->default_value(8), "number of threads")
    ("n_rays", po::value<long>()->default_value(10000000), "number of rays for primary emission")
    ("n_rays_int", po::value<long>()->default_value(10000000), "number of rays for intensity calc")
    ("n_rays_ir", po::value<long>()->default_value(10000000), "number of rays for IR emission")
    ("maxlevel", po::value<int>()->default_value(13), 
     "max grid refinement level")
    ("n_forced", po::value<int>()->default_value(1), 
     "number of forced scatterings")
    ("n_ir_forced", po::value<int>()->default_value(0), 
     "number of forced scatterings for IR emission")
    ("vis_ref", po::value<T_float>()->default_value(0.9e-6),
     "stellar radiation reference wavelength")
    ("ir_ref", po::value<T_float>()->default_value(50e-6),
     "dust emission reference wavelength")
    ("i_rr", po::value<T_float>()->default_value(1e-2),
     "intensity where Russian Roulette starts")
    ("abs_tol", po::value<T_float>()->default_value(1e-6), 
     "absolute ir equilibrium tolerance")
    ("rel_tol", po::value<T_float>()->default_value(1e-2),
     "relative ir equilibrium tolerance")
    ("rho_gr", po::value<T_float>()->default_value(8e3),
     "Ground density  (kg/m^3)")
    ("h2o_rh", po::value<T_float>()->default_value(1),
     "relative humidity of atmosphere")
    ("seed", po::value<int>()->default_value(42), "random number seed")
    ("do_ir", po::value<bool>()->default_value(true), "calculate IR emission")
    ("do_equilibrium", po::value<bool>()->default_value(true),
     "calculate dust self-absorption")
    ("help", "print this page")
    ;

  po::positional_options_description p;
  p.add("output-file", 1);

  po::variables_map opt;
  po::store(po::command_line_parser(argc, argv).
          options(desc).positional(p).run(), opt);
  po::notify(opt);    

  if (opt.count("help")) {
    cout << desc << "\n";
    return 1;
  }
  if(!opt.count("output-file")) {
    cout << "Must specify output file name on command line" << endl;
    return 1;
  }

  cout << "Running greenhouse model with parameters:\n";
  cout << "output-file = " << opt["output-file"].as<string>() << '\n'
       << "n_threads = " << opt["n_threads"].as<int>() << '\n'
       << "n_rays = " << opt["n_rays"].as<long>() << '\n'
       << "n_rays_int = " << opt["n_rays_int"].as<long>() << '\n'
       << "n_rays_ir = " << opt["n_rays_ir"].as<long>() << '\n'
       << "maxlevel = " << opt["maxlevel"].as<int>() << '\n'
       << "n_forced = " << opt["n_forced"].as<int>() << '\n'
       << "n_ir_forced = " << opt["n_ir_forced"].as<int>() << '\n'
       << "vis_ref = " << opt["vis_ref"].as<T_float>() << '\n'
       << "ir_ref = " << opt["ir_ref"].as<T_float>() << '\n'
       << "i_rr = " << opt["i_rr"].as<T_float>() << '\n'
       << "abs_tol = " << opt["abs_tol"].as<T_float>() << '\n'
       << "rel_tol = " << opt["rel_tol"].as<T_float>() << '\n'
       << "i_rr = " << opt["i_rr"].as<T_float>() << '\n'
       << "rho_gr = " << opt["rho_gr"].as<T_float>() << '\n'
       << "h2o_rh = " << opt["h2o_rh"].as<T_float>() << '\n'
       << "seed = " << opt["seed"].as<int>() << '\n'
       << "do_ir = " << opt["do_ir"].as<bool>() << '\n'
       << "do_equilibrium = " << opt["do_equilibrium"].as<bool>() << "\n\n";

  greenhouse_test ( opt );
}
